package com.ww.aiinterviewer.manager;

import cn.hutool.core.collection.CollUtil;
import com.volcengine.ark.runtime.model.bot.completion.chat.BotChatCompletionRequest;
import com.volcengine.ark.runtime.model.bot.completion.chat.BotChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChoice;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * AI调用工具类
 * @author wangwang
 * @create 2025/5/9 16:45
 */
@Service
public class AiManager {

    @Resource
    private ArkService arkService;

    /**
     * 聊天（只允许传入系统预设和用户预设）
     * @param systemPrompt 系统预设
     * @param userPrompt 用户预设
     * @author wangwang
     * @date 2025/5/9
     */
    public String doChat(String systemPrompt, String userPrompt) {
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(systemPrompt).build();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(userPrompt).build();
        messages.add(systemMessage);
        messages.add(userMessage);

        return doChat(messages);
    }

    /**
     * 更通用的方法，允许用户传入任意条消息列表
     * @param chatMessageList
     * @author wangwang
     * @date 2025/5/9
     */
    public String doChat(List<ChatMessage> chatMessageList) {

        // 单次调用
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                // 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
                .model("deepseek-v3-250324")
                .messages(chatMessageList)
                .build();

//        BotChatCompletionRequest chatCompletionRequest = BotChatCompletionRequest.builder()
//                .botId("bot-20250509150135-rdnp2")
//                .messages(chatMessageList)
//                .build();

        List<ChatCompletionChoice> choiceList = arkService.createChatCompletion(chatCompletionRequest)
                .getChoices();
//        List<ChatCompletionChoice> choiceList =  arkService.createBotChatCompletion(chatCompletionRequest).getChoices();
        if (CollUtil.isEmpty(choiceList)) {
            throw new RuntimeException("调用失败");
        }
        String content = (String) choiceList.get(0).getMessage().getContent();
        System.out.printf("AI 返回内容：%s\n", content);
        return content;
    }
}
